/**
 * @file ebex_mempool.c
 *
 * @date Oct 21, 2011
 * @author seth
 *
 * @brief This file is part of EBEX State Estimator, created for the EBEX project
 *
 * This software is copyright (C) 2011 Columbia University
 *
 * EBEX State Estimator is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * EBEX State Estimator is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with EBEX State Estimator; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */



#include <omp.h>
#include <ebex_mempool.h>

#define MEMORY_AREA_SIZE 4L

typedef void *memchunk_key_t;

typedef struct node_s
{
	struct node_s *left; /* Left subtree. */
	struct node_s *right; /* Right subtree. */
	int balance; /* Height (left) - height (right). */
	memchunk_key_t key; /* The key for this node. */
	void *data; /* Data stored at this node. */
} node_t;

typedef struct
{
	node_t *root;
} tree_t;

typedef struct atomic_free_s
{
	struct atomic_free_s *next;
} atomic_free_t;

typedef struct memarena_s
{
	struct memarena_s *next; /* The next memory area */
	struct memarena_s *prev; /* The previous memory area */
	size_t index; /* The current index into the "mem" array */
	size_t free; /* The number of free bytes in this memory area */
	unsigned int used; /* The number of atoms allocated from this area */
	unsigned char *mem; /* The mem array from which atoms get allocated. */
} memarena_t;

struct memchunk_s
{
	unsigned int num_mem_areas; /* The total number of memory areas */
	unsigned int num_unused_areas; /* The number of areas that may be deallocated */
	size_t atom_size; /* The size of an atom (used in mimic routines) */
	size_t area_size; /* The size of a memory area */
	memarena_t *mem_arena; /* The current memory area */
	memarena_t *mem_arenas; /* A list of all the memory areas owned by this chunk */
	memarena_t *free_mem_arena; /* Free area... which is about to be destroyed */
	atomic_free_t *free_atoms; /* Free atoms list */
	tree_t *mem_tree; /* Tree of memory areas sorted by memory address */
	long num_atoms_alloc; /* The number of allocated atoms (only used in mimic routines) */
};

/*
 * Private function prototypes.
 */
static node_t *node_new(memchunk_key_t key, void *data);
static void node_free(node_t *node);
static void node_delete(node_t *node);
static node_t *node_insert(node_t *node, memchunk_key_t key, void *data, bool *inserted);
static node_t *node_remove(node_t *node, memchunk_key_t key, void **removed_data);
static node_t *node_balance(node_t *node);
static node_t *node_remove_leftmost(node_t *node, node_t **leftmost);
static node_t *node_restore_left_balance(node_t *node, int old_balance);
static node_t *node_restore_right_balance(node_t *node, int old_balance);
static int node_height(node_t *node);
static node_t *node_rotate_left(node_t *node);
static node_t *node_rotate_right(node_t *node);

/*
 * Compilation constants.
 */
#define _NODE_BUFFER_NUM_INCR	16
#define _NODE_BUFFER_SIZE		1024

/*
 * Global variables.
 */
static node_t *node_free_list = NULL;
static int num_trees = 0; /* Count number of trees in use. */
static int buffer_num = -1;
static int num_buffers = 0;
static int num_used = _NODE_BUFFER_SIZE;
static node_t **node_buffers = NULL;

/*
 * Note that a single coarse thread lock is used when a number of
 * less coarse locks might be better.
 */
static omp_lock_t(node_buffer_lock);
static bool mem_chunk_openmp_initialised = false;

/*
 * This function must be called before any other functions is OpenMP
 * code is to be used.  Can be safely called when OpenMP code is not
 * being used, and can be safely called more than once.
 */
void mem_chunk_init_openmp(void)
{
	if (mem_chunk_openmp_initialised == false)
	{
		omp_init_lock(&node_buffer_lock);
		mem_chunk_openmp_initialised = true;
	}

	return;
}

/*
 * Private functions.
 */

/*
 * Deallocate all memory associated with buffers.  Should
 * never be called outside of a "node_buffer_lock"
 * block.
 */
static void destroy_buf(void)
{

	while (buffer_num >= 0)
	{
		e_free(node_buffers[buffer_num]);
		buffer_num--;
	}

	e_free(node_buffers);
	node_buffers = NULL;
	num_buffers = 0;
	num_used = _NODE_BUFFER_SIZE;

	node_free_list = NULL;

	return;
}

static node_t *node_new(memchunk_key_t key, void *data)
{
	node_t *node;

	omp_set_lock(&node_buffer_lock);
	/*
	 * Find an unused node.  Look for unused node in current buffer, then
	 * in the free node list, then allocate new buffer.
	 */
	if (num_used < _NODE_BUFFER_SIZE)
	{
		node = &(node_buffers[buffer_num][num_used]);
		num_used++;
	}
	else
	{
		if (node_free_list)
		{
			node = node_free_list;
			node_free_list = node->right;
		}
		else
		{
			buffer_num++;

			if (buffer_num == num_buffers)
			{
				num_buffers += _NODE_BUFFER_NUM_INCR;
				node_buffers = e_realloc(node_buffers, sizeof(node_t *)*num_buffers);
			}

			node_buffers[buffer_num] = e_malloc(sizeof(node_t) * _NODE_BUFFER_SIZE);

			if (!node_buffers[buffer_num])
				die("Unable to allocate memory.");

			node = node_buffers[buffer_num];
			num_used = 1;
		}
	}
	omp_unset_lock(&node_buffer_lock);

	node->balance = 0;
	node->left = NULL;
	node->right = NULL;
	node->key = key;
	node->data = data;

	return node;
}

static void node_free(node_t *node)
{
	omp_set_lock(&node_buffer_lock);
	node->right = node_free_list;
	node_free_list = node;
	omp_unset_lock(&node_buffer_lock);

	return;
}

static void node_delete(node_t *node)
{
	if (node)
	{
		node_delete(node->right);
		node_delete(node->left);
		node_free(node);
	}

	return;
}

/*
 * Systematically search tree with a search function which has the
 * same ordering as the tree.
 * This iterative version is much faster than the equivalent recursive version.
 */
static void *node_ordered_search(node_t *node, void *userdata)
{

	while (node)
	{
		if (userdata < (void *) ((memarena_t *) node->data)->mem)
			node = node->left;
		else if (userdata > (void *) &(((memarena_t *) node->data)->mem[((memarena_t *) node->data)->index]))
			node = node->right;
		else
			return node->data;
	}

	return NULL;
}

static node_t *node_insert(node_t *node, memchunk_key_t key, void *data, bool *inserted)
{
	int old_balance;

	if (!node)
	{
		*inserted = true;
		return node_new(key, data);
	}

	if (key < node->key)
	{
		if (node->left)
		{
			old_balance = node->left->balance;
			node->left = node_insert(node->left, key, data, inserted);

			if ((old_balance != node->left->balance) && node->left->balance)
				node->balance--;
		}
		else
		{
			*inserted = true;
			node->left = node_new(key, data);
			node->balance--;
		}
	}
	else if (key > node->key)
	{
		if (node->right)
		{
			old_balance = node->right->balance;
			node->right = node_insert(node->right, key, data, inserted);

			if ((old_balance != node->right->balance) && node->right->balance)
				node->balance++;
		}
		else
		{
			*inserted = true;
			node->right = node_new(key, data);
			node->balance++;
		}
	}
	else
	{ /* key == node->key */
		/*
		 *inserted = false;
		 */
		printf("WARNING: -- Replaced node -- (Key clash?)\n");

		node->data = data;
		return node;
	}

	if (*inserted && (node->balance < -1 || node->balance > 1))
		node = node_balance(node);

	return node;
}

static node_t *node_remove(node_t *node, memchunk_key_t key, void **removed_data)
{
	node_t *new_root;
	int old_balance;

	if (!node)
		return NULL;

	if (key < node->key)
	{
		if (node->left)
		{
			old_balance = node->left->balance;
			node->left = node_remove(node->left, key, removed_data);
			node = node_restore_left_balance(node, old_balance);
		}
	}
	else if (key > node->key)
	{
		if (node->right)
		{
			old_balance = node->right->balance;
			node->right = node_remove(node->right, key, removed_data);
			node = node_restore_right_balance(node, old_balance);
		}
	}
	else if (key == node->key)
	{
		node_t *removed_node;

		removed_node = node;

		if (!node->right)
		{
			node = node->left;
		}
		else
		{
			old_balance = node->right->balance;
			node->right = node_remove_leftmost(node->right, &new_root);
			new_root->left = node->left;
			new_root->right = node->right;
			new_root->balance = node->balance;
			node = node_restore_right_balance(new_root, old_balance);
		}

		*removed_data = removed_node->data;
		node_free(removed_node);
	}

	return node;
}

static node_t *node_balance(node_t *node)
{
	if (node->balance < -1)
	{
		if (node->left->balance > 0)
			node->left = node_rotate_left(node->left);
		node = node_rotate_right(node);
	}
	else if (node->balance > 1)
	{
		if (node->right->balance < 0)
			node->right = node_rotate_right(node->right);
		node = node_rotate_left(node);
	}

	return node;
}

static node_t *node_remove_leftmost(node_t *node, node_t **leftmost)
{
	int old_balance;

	if (!node->left)
	{
		*leftmost = node;
		return node->right;
	}

	old_balance = node->left->balance;
	node->left = node_remove_leftmost(node->left, leftmost);
	return node_restore_left_balance(node, old_balance);
}

static node_t *node_restore_left_balance(node_t *node, int old_balance)
{
	if ((!node->left) || ((node->left->balance != old_balance) && (node->left->balance == 0)))
	{
		node->balance++;
	}

	if (node->balance > 1)
		return node_balance(node);

	return node;
}

static node_t *node_restore_right_balance(node_t *node, int old_balance)
{
	if ((!node->right) || ((node->right->balance != old_balance) && (node->right->balance == 0)))
	{
		node->balance--;
	}

	if (node->balance < -1)
		return node_balance(node);

	return node;
}

static int node_height(node_t *node)
{
	int left_height;
	int right_height;

	if (!node)
		return 0;

	if (node->left)
		left_height = node_height(node->left);
	else
		left_height = 0;

	if (node->right)
		right_height = node_height(node->right);
	else
		right_height = 0;

	return MAX(left_height, right_height) + 1;
}

static node_t *node_rotate_left(node_t *node)
{
	node_t *right;
	int a_bal;
	int b_bal;

	right = node->right;

	node->right = right->left;
	right->left = node;

	a_bal = node->balance;
	b_bal = right->balance;

	if (b_bal <= 0)
	{
		if (a_bal >= 1)
			right->balance = b_bal - 1;
		else
			right->balance = a_bal + b_bal - 2;

		node->balance = a_bal - 1;
	}
	else
	{
		if (a_bal <= b_bal)
			right->balance = a_bal - 2;
		else
			right->balance = b_bal - 1;

		node->balance = a_bal - b_bal - 1;
	}

	return right;
}

static node_t *node_rotate_right(node_t *node)
{
	node_t *left;
	int a_bal;
	int b_bal;

	left = node->left;

	node->left = left->right;
	left->right = node;

	a_bal = node->balance;
	b_bal = left->balance;

	if (b_bal <= 0)
	{
		if (b_bal > a_bal)
			left->balance = b_bal + 1;
		else
			left->balance = a_bal + 2;

		node->balance = a_bal - b_bal + 1;
	}
	else
	{
		if (a_bal <= -1)
			left->balance = b_bal + 1;
		else
			left->balance = a_bal + b_bal + 2;

		node->balance = a_bal + 1;
	}

	return left;
}

static tree_t *tree_new(void)
{
	tree_t *tree;

	num_trees++;

	if (!(tree = e_malloc(sizeof(tree_t))))
		die("Unable to allocate memory.");

	tree->root = NULL;

	return tree;
}

static void delete(tree_t *tree)
{
	if (!tree)
		return;

	node_delete(tree->root);

	e_free(tree);

	num_trees--;

	omp_set_lock(&node_buffer_lock);
	if (num_trees == 0)
		destroy_buf();
	omp_unset_lock(&node_buffer_lock);

	return;
}

static bool insert(tree_t *tree, void *data)
{
	bool inserted = false;

	if (!data)
		die("Internal error: Trying to insert NULL data.");
	if (!tree)
		die("Internal error: Trying to insert into NULL tree.");

	tree->root = node_insert(tree->root, (memchunk_key_t) ((memarena_t *) data)->mem, data, &inserted);

	return inserted;
}

static void *remove_data(tree_t *tree, void *data)
{
	void *removed = NULL;

	if (!tree || !tree->root)
		return NULL;

	tree->root = node_remove(tree->root, (memchunk_key_t) ((memarena_t *) data)->mem, &removed);

	return removed;
}

static void *remove_key(tree_t *tree, memchunk_key_t key)
{
	void *removed = NULL;

	if (!tree || !tree->root)
		return NULL;

	tree->root = node_remove(tree->root, key, &removed);

	return removed;
}

static void *ordered_search(tree_t *tree, void *userdata)
{
	if (!tree || !tree->root)
		return NULL;

	return node_ordered_search(tree->root, userdata);
}

static memchunk_t *mem_chunk_new_internal(size_t atom_size, unsigned int num_atoms)
{
	memchunk_t *mem_chunk;

	/*
	 * Ensure that we don't misalign allocated memory for the user.
	 * This also ensures that the minimum atom_size is okay for the
	 * FreeAtom list.
	 */
	if (atom_size % MEMORY_ALIGN_SIZE > 0)
	{
		atom_size += MEMORY_ALIGN_SIZE - (atom_size % MEMORY_ALIGN_SIZE);
		printf("DEBUG: modified Memchunk atomic size.\n");
	}

	if (!(mem_chunk = (memchunk_t *) e_malloc(sizeof(memchunk_t))))
		die("Unable to allocate memory.");

	mem_chunk->num_mem_areas = 0;
	mem_chunk->num_unused_areas = 0;
	mem_chunk->mem_arena = NULL;
	mem_chunk->free_mem_arena = NULL;
	mem_chunk->free_atoms = NULL;
	mem_chunk->mem_arenas = NULL;
	mem_chunk->atom_size = atom_size;
	mem_chunk->area_size = atom_size * num_atoms;
	mem_chunk->mem_tree = NULL;

	return mem_chunk;
}

/*
 * Return true is the memory chunk is empty.
 */bool mem_chunk_isempty(memchunk_t *mem_chunk)
{

	if (!mem_chunk)
		die("Null pointer to mem_chunk passed.");

	return (bool) (mem_chunk->num_mem_areas == mem_chunk->num_unused_areas);
}


memchunk_t *mem_chunk_new(size_t atom_size, unsigned int num_atoms)
{
	memchunk_t *mem_chunk;

	if (atom_size < 1)
		die("Passed atom size is < 1 byte.");
	if (num_atoms < 1)
		die("Passed number of atoms is < 1.");

	mem_chunk = mem_chunk_new_internal(atom_size, num_atoms);
	mem_chunk->mem_tree = tree_new();

	return mem_chunk;
}

void mem_chunk_destroy(memchunk_t *mem_chunk)
{
	memarena_t *mem_areas;
	memarena_t *temp_area;

	if (!mem_chunk)
		die("Null pointer to mem_chunk passed.");

	mem_areas = mem_chunk->mem_arenas;
	while (mem_areas)
	{
		temp_area = mem_areas;
		mem_areas = mem_areas->next;
		e_free(temp_area);
	}

	delete(mem_chunk->mem_tree);

	e_free(mem_chunk);

	return;
}

void *mem_chunk_alloc(memchunk_t *mem_chunk)
{
	memarena_t *temp_area;
	void *mem;

	if (!mem_chunk)
		die("Null pointer to mem_chunk passed.");

	while (mem_chunk->free_atoms)
	{
		/* Get the first piece of memory on the "free_atoms" list.
		 * We can go ahead and destroy the list node we used to keep
		 *  track of it with and to update the "free_atoms" list to
		 *  point to its next element.
		 */
		mem = mem_chunk->free_atoms;
		mem_chunk->free_atoms = mem_chunk->free_atoms->next;

		/* Determine which area this piece of memory is allocated from */
		temp_area = ordered_search(mem_chunk->mem_tree, mem);

		/* If the area is unused, then it may be destroyed.
		 * We check to see if all of the segments on the free list that
		 *  reference this area have been removed. This occurs when
		 *  the amount of free memory is less than the size that can be allocated.
		 * If the chunk should be freed, then we place it in the "free_mem_area".
		 * This is so we make sure not to free the memory area here and then
		 *  allocate it again a few lines down.
		 * If we don't allocate a chunk a few lines down then the "free_mem_area"
		 *  will be freed.
		 * If there is already a "free_mem_area" then we'll just free this memory area.
		 */
		if (temp_area->used == 0)
		{
			/* Update the "free" memory available in that area */
			temp_area->free += mem_chunk->atom_size;

			if (temp_area->free == mem_chunk->area_size)
			{
				if (temp_area == mem_chunk->mem_arena)
					mem_chunk->mem_arena = NULL;

				if (mem_chunk->free_mem_arena)
				{
					mem_chunk->num_mem_areas--;

					if (temp_area->next)
						temp_area->next->prev = temp_area->prev;
					if (temp_area->prev)
						temp_area->prev->next = temp_area->next;
					if (temp_area == mem_chunk->mem_arenas)
						mem_chunk->mem_arenas = mem_chunk->mem_arenas->next;

					if (mem_chunk->mem_tree)
					{
						if (!remove_data(mem_chunk->mem_tree, temp_area))
							die("Unable to find temp_area.");
					}

					e_free(temp_area);
				}
				else
					mem_chunk->free_mem_arena = temp_area;

				mem_chunk->num_unused_areas--;
			}
		}
		else
		{
			/* Update the number of allocated atoms count.
			 */
			temp_area->used++;

			/* The area is still in use...return the memory
			 */

			return mem;
		}
	}

	/* If there isn't a current memory area or the current memory area is out of
	 * space then allocate a new memory area. We'll first check and see if we can
	 * use the "free_mem_area".  Otherwise we'll just malloc the memory area.
	 */
	if ((!mem_chunk->mem_arena) || ((mem_chunk->mem_arena->index + mem_chunk->atom_size) > mem_chunk->area_size))
	{
		if (mem_chunk->free_mem_arena)
		{
			mem_chunk->mem_arena = mem_chunk->free_mem_arena;
			mem_chunk->free_mem_arena = NULL;
		}
		else
		{
			mem_chunk->mem_arena = (memarena_t*) e_malloc(
					sizeof(memarena_t) + MEMORY_ALIGN_SIZE - (sizeof(memarena_t) % MEMORY_ALIGN_SIZE)
							+ mem_chunk->area_size);

			if (!mem_chunk->mem_arena)
				die("Unable to allocate memory.");

			mem_chunk->mem_arena->mem = ((unsigned char*) (mem_chunk->mem_arena) + sizeof(memarena_t)
					+ MEMORY_ALIGN_SIZE - (sizeof(memarena_t) % MEMORY_ALIGN_SIZE));

			mem_chunk->num_mem_areas++;
			mem_chunk->mem_arena->next = mem_chunk->mem_arenas;
			mem_chunk->mem_arena->prev = NULL;

			if (mem_chunk->mem_arenas)
				mem_chunk->mem_arenas->prev = mem_chunk->mem_arena;
			mem_chunk->mem_arenas = mem_chunk->mem_arena;

			if (mem_chunk->mem_tree)
				insert(mem_chunk->mem_tree, mem_chunk->mem_arena);
		}

		mem_chunk->mem_arena->index = 0;
		mem_chunk->mem_arena->free = mem_chunk->area_size;
		mem_chunk->mem_arena->used = 0;
	}

	/*
	 * Get the memory and modify the state variables appropriately.
	 */
	mem = (void*) &(mem_chunk->mem_arena->mem[mem_chunk->mem_arena->index]);
	mem_chunk->mem_arena->index += mem_chunk->atom_size;
	mem_chunk->mem_arena->free -= mem_chunk->atom_size;
	mem_chunk->mem_arena->used++;

	return mem;
}

void mem_chunk_free(memchunk_t *mem_chunk, void *mem)
{
	memarena_t *temp_area;
	atomic_free_t *free_atom;

	if (!mem_chunk)
		die("Null pointer to mem_chunk passed.");
	if (!mem_chunk->mem_tree)
		die("Memchunk passed has no freeable atoms.");
	if (!mem)
		die("NULL pointer passed.");

	/*
	 * Place the memory on the "free_atoms" list.
	 */
	free_atom = (atomic_free_t*) mem;
	free_atom->next = mem_chunk->free_atoms;
	mem_chunk->free_atoms = free_atom;

	if (!(temp_area = ordered_search(mem_chunk->mem_tree, mem)))
		die("Unable to find temp_area.");

	temp_area->used--;

	if (temp_area->used == 0)
	{
		mem_chunk->num_unused_areas++;
	}

	return;
}

/* This doesn't free the free_area if there is one */
void mem_chunk_clean(memchunk_t *mem_chunk)
{
	memarena_t *mem_area;
	atomic_free_t *prev_free_atom;
	atomic_free_t *temp_free_atom;
	void *mem;

	if (!mem_chunk)
		die("Null pointer to mem_chunk passed.");
	if (!mem_chunk->mem_tree)
		die("Memchunk passed has no freeable atoms.");

	if (mem_chunk->mem_tree)
	{
		prev_free_atom = NULL;
		temp_free_atom = mem_chunk->free_atoms;

		while (temp_free_atom)
		{
			mem = (void*) temp_free_atom;

			if (!(mem_area = ordered_search(mem_chunk->mem_tree, mem)))
				die("mem_area not found.");

			/*
			 * If this memory area is unused, delete the area, list node and decrement the free mem.
			 */
			if (mem_area->used == 0)
			{
				if (prev_free_atom)
					prev_free_atom->next = temp_free_atom->next;
				else
					mem_chunk->free_atoms = temp_free_atom->next;
				temp_free_atom = temp_free_atom->next;

				mem_area->free += mem_chunk->atom_size;
				if (mem_area->free == mem_chunk->area_size)
				{
					mem_chunk->num_mem_areas--;
					mem_chunk->num_unused_areas--;

					if (mem_area->next)
						mem_area->next->prev = mem_area->prev;
					if (mem_area->prev)
						mem_area->prev->next = mem_area->next;
					if (mem_area == mem_chunk->mem_arenas)
						mem_chunk->mem_arenas = mem_chunk->mem_arenas->next;
					if (mem_area == mem_chunk->mem_arena)
						mem_chunk->mem_arena = NULL;

					if (mem_chunk->mem_tree)
					{
						if (!remove_key(mem_chunk->mem_tree, (memchunk_key_t) mem_area))
							die("mem_area not found.");
					}

					e_free(mem_area);
				}
			}
			else
			{
				prev_free_atom = temp_free_atom;
				temp_free_atom = temp_free_atom->next;
			}
		}
	}

	return;
}

void mem_chunk_reset(memchunk_t *mem_chunk)
{
	memarena_t *mem_areas;
	memarena_t *temp_area;

	if (!mem_chunk)
		die("Null pointer to mem_chunk passed.");

	mem_areas = mem_chunk->mem_arenas;
	mem_chunk->num_mem_areas = 0;
	mem_chunk->mem_arenas = NULL;
	mem_chunk->mem_arena = NULL;

	while (mem_areas)
	{
		temp_area = mem_areas;
		mem_areas = mem_areas->next;
		e_free(temp_area);
	}

	mem_chunk->free_atoms = NULL;

	if (mem_chunk->mem_tree)
	{
		delete(mem_chunk->mem_tree);
		mem_chunk->mem_tree = tree_new();
	}

	return;
}

